import numpy as np
from scipy.stats import norm, t, chi2, f
from scipy.integrate import quad, dblquad
from sympy import Rational as R
from sympy.utilities.iterables import variations as permutations
from sympy.utilities.iterables import subsets as combinations
from scipy.stats import rv_continuous
def subSet(A, condition):
    B=set()
    for x in A:
        if condition(x):
            B.add(x)
    return B
def P(A, S):
    n = len(S)
    m = len(A)
    return R(m, n)
def areaBetween(a, b, lowfunc, upfunc):
    one = lambda x, y: 1
    area, _ =dblquad(one, a, b, lowfunc, upfunc)
    return area
def condP(A, B):
    #return len(A & B) / len(B)
    return R(len(A & B), len(B))
def totalProb(likelihood, prioProbs):
    return (likelihood*prioProbs).sum()
def bayes(liklihood, prioProbs, i):
    temp=liklihood*prioProbs
    total=temp.sum()
    return temp[i-1]/total
def margDist(XY):
    X=XY.sum(axis=1)
    Y=XY.sum(axis=0)
    return X.reshape(X.size, 1), Y.reshape(1, Y.size)
def condDist(XY):
    X, Y=margDist(XY)
    Y_X=XY/X
    X_Y=XY/Y
    return X_Y, Y_X
def distCalcu(Y, P):
    Y1=np.unique(Y)
    n=len(Y1)
    P1=np.zeros(n)
    for i in range(n):
        P1[i]=np.sum(P[np.where(Y==Y1[i])])
    return (Y1, P1)
def independent(Pxy):
    Px, Py=margDist(Pxy)
    PxPy=Px*Py
    if PxPy.dtype==float:
        return (abs(PxPy-Pxy)<1e-8).all()
    return (PxPy==Pxy).all()
def appointment(a, b, c, d, e):
    f=lambda y, x: 1/(b-a)/(d-c) if (a<x)&(x<b)&(c<y)&(y<d) else 0
    g1=lambda x: x-e
    g2=lambda x: x+e
    p, _=dblquad(f, a, b, g1, g2)
    return p
def expect(P, Xv=None, Yv=None, func=lambda x, y: x):
    if type(Xv) != type(None):
        Xv=Xv.reshape(Xv.size,1)
    if type(Yv) != type(None):
        Yv=Yv.reshape(1, Yv.size)
    mean=(P*func(Xv,Yv)).sum()
    return mean
class basecont(rv_continuous):
    def var(self):
        moment1=self.expect()
        moment2=self.expect(func=lambda x: x**2)
        return moment2-moment1**2
def variance(P, X):
    g=lambda x, y:x**2
    Ex2=expect(P, Xv=X, func=g)
    Ex=expect(P, X)
    return Ex2-Ex**2
def expectcont(pdf, func=lambda x:x):
    f=lambda x:pdf(x)*func(x)
    mean, _=quad(f, -np.infty, np.infty)
    return mean
def expectcont2(pdf, func):
    gf=lambda y, x:pdf(y, x)*func(y, x)
    mean, _=dblquad(gf, -np.infty, np.infty,
                        -np.infty, np.infty)
    return mean
def cov(Exy, Ex, Ey):
    return Exy-Ex*Ey
def rhoxy(Exy, Ex, Ey, sigmax, sigmay):
    return cov(Exy, Ex, Ey)/sigmax/sigmay
def Fn(x, sample, m):
    n=sample.size
    f, b=np.histogram(sample, bins=m)
    f=f/n
    for k in range(m-1):
        f[k+1]+=f[k]
    y=np.zeros(x.size)
    for i in range(1, m):
        d=np.where((x>b[i-1])&(x<=b[i]))
        y[d]=f[i-1]
    d=np.where(x>b[m-1])
    y[d]=1
    return y
def muBounds(mean, d, confidence, df=0):
    if df==0:
        a, b=norm.interval(confidence)
    else:
        a, b=t.interval(confidence, df)
    return mean+a*d, mean+b*d
def muBound(mean, s, confidence, df=0, low=True):
    alpha=1-confidence
    if df==0:
        b=norm.isf(alpha)
    else:
        b=t.isf(alpha, df)
    if low:
        return mean-b*s
    return mean+b*s
def sigma2Bounds(d, df, confidence):
    a, b=chi2.interval(confidence, df)
    return d/b, d/a
def sigma2Bound(d, df, confidence, low=True):
    alpha=1-confidence
    if low:
        b=chi2.isf(alpha, df)
    else:
        b=chi2.ppf(alpha, df)
    return d/b
def sigma2RatioBounds(d, dfn, dfd, confidence):
    (a,b)=f.interval(confidence, dfn, dfd)
    return d/b, d/a
def sigma2RatioBound(d, dfn, dfd, confidence, low=True):
    alpha=1-confidence
    if low:
        a=f.isf(alpha, dfn, dfd)
    else:
        a=f.ppf(alpha, dfn, dfd)
    return d/a
def ztest(Z, alternative='two-sided'):
    if alternative=='greater':
        p=norm.sf(Z)
    else:
        if alternative=='less':
            p=norm.cdf(Z)
        else:
            p=2*norm.sf(abs(Z))
    return p
def ttest(T, df, alternative='two-sided'):
    if alternative=='greater':
        p=t.sf(T, df)
    else:
        if alternative=='less':
            p=t.cdf(T, df)
        else:
            p=2*t.sf(abs(T), df)
    return p
def chi2test(X, df, alternative='two-sided'):
    if alternative=='greater':
        p=chi2.sf(X, df)
    else:
        if alternative=='less':
            p=chi2.cdf(X, df)
        else:
            if X>df:
                p=2*chi2.sf(X, df)
            else:
                p=2*chi2.cdf(X, df)
    return p
def ftest(F, dfn, dfd, alternative='two-sided'):
    if alternative=='greater':
        p=f.sf(F, dfn, dfd)
    else:
        if alternative=='less':
            p=f.cdf(F, dfn, dfd)
        else:
            if F>dfd/(dfd-2):
                p=2*f.sf(F, dfn, dfd)
            else:
                p=2*f.cdf(F, dfn, dfd)
    return p
def sfeDecompose(X):
    s=X.shape[0]
    n=np.array([X[i].size for i in range(s)])
    nt=n.sum()
    X_bar=np.array([X[i].mean() for i in range(s)])
    Xt_bar=(X_bar*n).sum()/nt
    ST=np.sum([((X[i]-Xt_bar)**2).sum() for i in range(s)])
    SA=(n*(X_bar**2)).sum()-nt*(Xt_bar**2) 
    SE=ST-SA
    return (n, s, X_bar, Xt_bar, ST, SA, SE)
def sfeTest(n, s, SA, SE, alpha):
    nt=n.sum()
    k=f.isf(alpha, s-1, nt-s)
    F=(nt-s)/(s-1)*SA/SE
    return F<k
def sfeEstimat(accept, n, s, X_bar, Xt_bar, ST, SE, alpha):
    ans=[]
    nt=n.sum()
    (a, b)=sigma2Bounds(SE, nt-s, 1-alpha)
    ans.append((a, b))
    if accept: 
        d=np.sqrt(ST/(nt-1)/nt)
        (a, b)=muBounds(Xt_bar, d, 1-alpha, nt-1)
        ans.append((a, b))
    else:
        for i in range(s):
            for j in range(i+1, s):
                mean=X_bar[i]-X_bar[j]
                S_E=SE/(nt-s)
                d=np.sqrt(S_E*(1/n[i]+1/n[j]))
                (a, b)=muBounds(mean, d, 1-alpha, nt-s)
                ans.append((a, b))
    return np.array(ans)
def sfeVarianceAnalysis(X,alpha):
    (n, s, X_bar, Xt_bar, ST, SA, SE)=sfeDecompose(X)
    accept=sfeTest(n, s, SA, SE, alpha)
    ans=sfeEstimat(accept, n, s, X_bar, Xt_bar, ST, SE, alpha)
    print('显著水平%.3f下H0为%s'%(alpha,accept))
    print('置信水平%.3f下：'%(1-alpha))
    print('sigma^2的置信区间为(%.6f, %.6f)'
          %(ans[0][0], ans[0][1]))
    if accept:
        print('mu的置信区间为(%.4f, %.4f)'
              %(ans[1][0], ans[1][1]))
    else:
        k=0
        for i in range(s):
            for j in range(i+1, s):
                print('mu%d-mu%d的置信区间为(%.4f, %.4f)'
                      %(i+1, j+1, ans[1+k][0], ans[1+k][1]))
                k=k+1
def dfeVarianceAnalysis(X, alpha):
    (r,s,t)=X.shape
    X_bar=X.mean(axis=2)
    Xi_bar=X_bar.mean(axis=1).reshape(r,1)
    Xj_bar=X_bar.mean(axis=0).reshape(1,s)
    Xt_bar=X.mean()
    ST=((X-Xt_bar)**2).sum()
    SA=s*t*((Xi_bar-Xt_bar)**2).sum()
    SB=r*t*((Xj_bar-Xt_bar)**2).sum()
    SAB=t*((X_bar-Xi_bar-Xj_bar+Xt_bar)**2).sum()
    SE=ST-SA-SB-SAB
    F1=(r*s*(t-1))*SA/SE/(r-1)
    k1=f.isf(alpha,r-1,r*s*(t-1))
    F2=(r*s*(t-1))*SB/SE/(s-1)
    k2=f.isf(alpha,s-1,r*s*(t-1))
    F3=(r*s*(t-1)*SAB)/((r-1)*(s-1)*SE)
    k3=f.isf(alpha,(r-1)*(s-1), r*s*(t-1))
    return np.array([F1<k1,F2<k2,F3<k3])
def dfeVarianceAnalysis1(X, alpha):
    r,s=X.shape
    Xi_bar=X.mean(axis=1).reshape(r, 1)
    Xj_bar=X.mean(axis=0).reshape(1, s)
    Xt_bar=X.mean()
    ST=((X-Xt_bar)**2).sum()
    SA=s*((Xi_bar-Xt_bar)**2).sum()
    SB=r*((Xj_bar-Xt_bar)**2).sum()
    SE=ST-SA-SB
    k1=f.isf(alpha, r-1, (r-1)*(s-1))
    k2=f.isf(alpha, s-1, (r-1)*(s-1))
    F1=(s-1)*SA/SE
    F2=(r-1)*SB/SE
    ans=np.array([F1<k1, F2<k2])
    return ans
def controlright(a, b, s, y1, alpha):
    Z=norm.ppf(alpha)
    c=y1-b
    dy=Z*s
    return (c-dy)/a
def controlleft(a, b, s, y1, alpha):
    Z=norm.isf(alpha)
    c=y1-b
    dy=Z*s
    return (c-dy)/a
def controlbi(a, b, s, y1, y2, alpha):
    Z1,Z2=norm.interval(1-alpha)
    c1=y1-b
    c2=y2-b
    dy1=Z1*s
    dy2=Z2*s
    p1=(c1-dy1)/a
    p2=(c2-dy2)/a
    if p2<p1:
        (p1,p2)=(p2,p1)
    return (p1, p2)